import datetime

from langchain.tools import tool, StructuredTool
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain.agents import initialize_agent, AgentType

from config.load_key import load_key

llm = ChatOpenAI(
    model="Qwen/Qwen2.5-7B-Instruct",
    api_key=load_key("siliconflow_api_key"),
    base_url=load_key("siliconflow_base_url"),
)


# 注解方式的工具
@tool("get_current_date")
def get_current_date():
    """获取今天日期"""
    return datetime.datetime.today().strftime("%Y-%m-%d")


@tool(description="获取某个城市的天气")
def get_city_weather(city: str):
    """获取某个城市的天气
    Args:
        city: 具体城市
    """
    return "城市" + city + "，今天天气不错"


# StructuredTool.from_function 方式的工具
def bad_weather_tool(city: str):
    """获取某个城市的天气
    Args:
        city: 具体城市
    """
    return "城市" + city + "，今天天气不太好"


weatherTool = StructuredTool.from_function(func=bad_weather_tool, description="获取某个城市的天气",
                                           name="bad_weather_tool")

# chain 方式的工具
prompt = ChatPromptTemplate.from_messages([("human", "你好，请用下面这种语言回答我的问题 {language}.")])

parser = StrOutputParser()

chain = prompt | llm | parser

as_tool = chain.as_tool(name="translatetool", description="翻译任务")

# 初始化代理
# AgentType 参考 https://blog.csdn.net/u013172930/article/details/147645558
agent = initialize_agent(
    tools=[get_city_weather, weatherTool, as_tool],
    llm=llm,
    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

#query = "今天是几号？今天北京天气怎么样？ 用中文回答一遍，再用英文回答一遍"
query = "北京天气怎么样？用中文回答一遍，再用英文回答一遍"
response = agent.invoke({"input": query})
print(response)

